#include <rclcpp/rclcpp.hpp>
#include "performance_test_base.h"
#include <thread>
#include <vector>
#include <atomic>

namespace trace_api = opentelemetry::trace;
namespace nostd = opentelemetry::nostd;

class ConcurrentTestNode : public rclcpp::Node {
public:
    ConcurrentTestNode() : Node("concurrent_test_node") {
        // Initialize OpenTelemetry
        PerformanceTestBase::InitTracer();
        tracer_ = trace_api::Provider::GetTracerProvider()->GetTracer("concurrent-test");
        
        RCLCPP_INFO(this->get_logger(), "Concurrent Span Test started");
        RCLCPP_INFO(this->get_logger(), "Testing concurrent span creation performance...");
        
        runConcurrentSpanTest();
        
        RCLCPP_INFO(this->get_logger(), "Concurrent span test completed!");
    }

private:
    void runConcurrentSpanTest() {
        RCLCPP_INFO(this->get_logger(), "Running concurrent span test...");
        
        PerformanceStats stats;
        stats.reset();
        
        stats.memory_before_kb = PerformanceTestBase::getCurrentMemoryKB();
        auto cpu_before = PerformanceTestBase::getCurrentCPUTime();
        stats.cpu_time_user_before = cpu_before.first;
        stats.cpu_time_system_before = cpu_before.second;
        
        const size_t num_threads = 4;
        const size_t spans_per_thread = 2500;
        std::vector<std::thread> threads;
        std::atomic<size_t> total_operations{0};
        
        auto start_time = std::chrono::high_resolution_clock::now();
        
        for (size_t t = 0; t < num_threads; ++t) {
            threads.emplace_back([this, spans_per_thread, t, &total_operations]() {
                for (size_t i = 0; i < spans_per_thread; ++i) {
                    auto span = tracer_->StartSpan("concurrent_span", {
                        {"thread_id", static_cast<int64_t>(t)},
                        {"span_id", static_cast<int64_t>(i)}
                    });
                    
                    span->SetAttribute("thread.data", "thread_" + std::to_string(t) + "_span_" + std::to_string(i));
                    span->End();
                    
                    total_operations++;
                }
            });
        }
        
        for (auto& thread : threads) {
            thread.join();
        }
        
        auto end_time = std::chrono::high_resolution_clock::now();
        stats.total_time = std::chrono::duration_cast<std::chrono::nanoseconds>(end_time - start_time);
        stats.operation_count = total_operations.load();
        
        stats.memory_after_kb = PerformanceTestBase::getCurrentMemoryKB();
        auto cpu_after = PerformanceTestBase::getCurrentCPUTime();
        stats.cpu_time_user_after = cpu_after.first;
        stats.cpu_time_system_after = cpu_after.second;
        
        stats.calculateAverages();
        PerformanceTestBase::printStats("Concurrent Span Creation", stats);
    }

    nostd::shared_ptr<trace_api::Tracer> tracer_;
};

int main(int argc, char* argv[]) {
    rclcpp::init(argc, argv);
    
    auto node = std::make_shared<ConcurrentTestNode>();
    
    RCLCPP_INFO(node->get_logger(), "Starting Concurrent Span Performance Test...");
    RCLCPP_INFO(node->get_logger(), "This will test OpenTelemetry's concurrent span creation performance");
    
    rclcpp::spin(node);
    rclcpp::shutdown();
    
    return 0;
}
